#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image.h"
#include "stb_image_write.h"
#include <iostream>
#include <cuda_runtime.h>

// CUDA kernel function for RGB to Grayscale conversion
__global__ void rgbToGray(unsigned char* input, unsigned char* output, int width, int height) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= width || y >= height) return;

    int rgbIndex = (y * width + x) * 3;  // Index for RGB image
    int grayIndex = y * width + x;  // Index for grayscale image

    unsigned char r = input[rgbIndex];
    unsigned char g = input[rgbIndex + 1];
    unsigned char b = input[rgbIndex + 2];

    // Use luminosity method to convert RGB to Grayscale
    output[grayIndex] = 0.21f * r + 0.72f * g + 0.07f * b;
}

void checkCudaError(cudaError_t err, const char* message) {
    if (err != cudaSuccess) {
        std::cerr << "Error: " << message << " (" << cudaGetErrorString(err) << ")" << std::endl;
        exit(EXIT_FAILURE);
    }
}

int main() {
    int width, height, channels;

    // Load input image using stb_image
    unsigned char* h_input = stbi_load("input.jpg", &width, &height, &channels, 3);
    if (!h_input) {
        std::cerr << "Error: Could not open or find the image." << std::endl;
        return -1;
    }

    // Allocate host memory for output (grayscale image)
    unsigned char* h_output = new unsigned char[width * height];

    // Allocate device memory
    unsigned char* d_input;
    unsigned char* d_output;
    checkCudaError(cudaMalloc(&d_input, width * height * 3), "cudaMalloc d_input");
    checkCudaError(cudaMalloc(&d_output, width * height), "cudaMalloc d_output");

    // Copy input data from host to device
    checkCudaError(cudaMemcpy(d_input, h_input, width * height * 3, cudaMemcpyHostToDevice), "cudaMemcpy h_input to d_input");

    // Define grid and block dimensions
    dim3 blockDim(16, 16);
    dim3 gridDim((width + blockDim.x - 1) / blockDim.x, (height + blockDim.y - 1) / blockDim.y);

    // Launch the CUDA kernel
    rgbToGray<<<gridDim, blockDim>>>(d_input, d_output, width, height);

    // Check for any kernel launch errors
    checkCudaError(cudaGetLastError(), "Kernel launch failed");

    // Copy output data from device to host
    checkCudaError(cudaMemcpy(h_output, d_output, width * height, cudaMemcpyDeviceToHost), "cudaMemcpy d_output to h_output");

    // Save output image using stb_image_write
    if (!stbi_write_jpg("output.jpg", width, height, 1, h_output, 100)) {
        std::cerr << "Error: Could not write the output image." << std::endl;
    }

    // Free device memory
    checkCudaError(cudaFree(d_input), "cudaFree d_input");
    checkCudaError(cudaFree(d_output), "cudaFree d_output");

    // Free host memory
    stbi_image_free(h_input);
    delete[] h_output;

    std::cout << "Image processing completed." << std::endl;

    return 0;
}
